from autogen_agentchat.agents import AssistantAgent
from autogen_core.models import UserMessage
from autogen_ext.models.ollama import OllamaChatCompletionClient
from autogen_agentchat.messages import MultiModalMessage
from autogen_core import Image as AGImage
from PIL import Image
from autogen_core import CancellationToken
# 测试ollama客户端
ollama_model_client = OllamaChatCompletionClient(model="qwen2.5vl:7b",
                                                 model_info={
                                                     "vision": True,
                                                     "function_calling": True,
                                                     "family": "qwen2.5",
                                                     "structured_output": True,
                                                     "json_output": True,
                                                 })




async def test_run_chat():
    print("*" * 50)
    print("测试client")
    pil_image = Image.open(r"../0250614220809.png")
    img = AGImage(pil_image)
    # response = await agent.on_messages([multi_modal_message],CancellationToken())
    # print(response.chat_message.content)
    # 尝试使用create方法
    response = await ollama_model_client.create([
        UserMessage(content=["描述这张照片"], source="user"),
        UserMessage(content=[img], source="user"),
    ],extra_create_args={
        "think": False
    })
    print(response)

async def test_run_assistant():
    print("*"*50)
    print("测试assistant")
    assistant = AssistantAgent(name="assistant", model_client=ollama_model_client)
    pil_image = Image.open(r"../0250614220809.png")
    img = AGImage(pil_image)
    print(img)
    multi_modal_message = MultiModalMessage(content=["描述这张图片?", img], source="User")
    response = await assistant.on_messages([multi_modal_message],CancellationToken())
    print(response.chat_message.content)
async def combined_test():
    print("*"*50)
    print("测试combined")
    await test_run_chat()
    await test_run_assistant()
if __name__ == "__main__":
    import asyncio
    # 运行测试
    asyncio.run(combined_test())


